import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import random_split
from torch.optim import lr_scheduler
from torchvision import datasets, models, transforms
import numpy as np
import torch.utils.data as Data
from torchvision.datasets import ImageFolder
import copy

model_res = models.resnet18(pretrained=True)
#model_res = models.resnet152(pretrained=True)

'''
resnet152不方便训练，这里先使用resnet18
'''

data_transform = transforms.Compose([
    transforms.Resize(size=(224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

batch_size = 128
train_num = 10000
valid_num = 1000

dataset = ImageFolder("/mnt/data",
                      transform=data_transform)
f=open('/mnt/indices.txt',"w")
print(dataset.class_to_idx,file=f)


temp_dataset, static_dataset = random_split(dataset, [11000, len(dataset)-11000])
train_dataset, valid_dataset = random_split(temp_dataset, [10000, 1000])

train_loader = Data.DataLoader(
    dataset=train_dataset, batch_size=batch_size, shuffle=True)
valid_loader = Data.DataLoader(
    dataset=valid_dataset, batch_size=batch_size, shuffle=True)

num_ftrs = model_res.fc.in_features
model_res.fc = nn.Linear(num_ftrs, 158).cuda()

criterion = nn.CrossEntropyLoss().cuda()
optimizer_res = optim.SGD(model_res.fc.parameters(), lr=0.01, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_res, step_size=7, gamma=0.1)
model_res = model_res.cuda()
dataloader = {
    'train': train_loader,
    'val': valid_loader
}

dataset_size = {
    'train': train_num,
    'val': valid_num
}


def train_model(model, criterion, optimizer, scheduler, num_epochs=5):
    best_model = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    # set GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    for epoch in range(num_epochs):
        print('Epoch %d: ' % epoch)
        print('-' * 10)

        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                model.train()
            else:
                model.eval()

            running_loss = 0.0
            running_corrects = 0
            i=0
            for inputs, labels in dataloader[phase]:
                i+=1
                print("%f"%(batch_size*i/10000))
                print('training!')
                inputs = inputs.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, pred = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(pred == labels.data)

            epoch_loss = running_loss / dataset_size[phase]
            epoch_acc = running_corrects.double() / dataset_size[phase]
            print('%s Loss = %.4f, Acc = %.4f' %
                  (phase, epoch_loss, epoch_acc))

            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model = copy.deepcopy(model.state_dict())

    print('Completed training!')
    print('best_acc = %.4f' % best_acc)
    model.load_state_dict(best_model)
    return model


model_final = train_model(model_res, criterion,
                          optimizer_res, exp_lr_scheduler)

torch.save(model_final,'result.pkl')
